package com.ideaaedi.zoo.diy.artifact.dynamic.datasource.datasource;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.ideaaedi.commonds.exec.Executor;
import com.ideaaedi.commonds.function.NoArgConsumer;
import com.ideaaedi.commonds.function.NoArgFunction;

import java.util.function.Consumer;
import java.util.function.Function;

/**
 * 数据源切换器
 *
 * @author <font size = "20" color = "#3CAA3C"><a href="https://gitee.com/JustryDeng">JustryDeng</a></font> <img
 * src="https://gitee.com/JustryDeng/shared-files/raw/master/JustryDeng/avatar.jpg" />
 * @since 1.0.0
 */
public final class DatasourceSwitcher implements Executor {
    
    private final String datasourceName;
    
    private DatasourceSwitcher(String datasourceName) {
        this.datasourceName = datasourceName;
    }
    
    /**
     * 切换至指定的数据源，然后在该数据源内执行相关逻辑
     */
    public static DatasourceSwitcher switchTo(String datasourceName) {
        return new DatasourceSwitcher(datasourceName);
    }
    
    @Override
    public <P, R> R exec(Function<P, R> function, P param) {
        try {
            DynamicDataSourceContextHolder.push(datasourceName);
            return function.apply(param);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
    
    @Override
    public <R> R exec(NoArgFunction<R> function) {
        try {
            DynamicDataSourceContextHolder.push(datasourceName);
            return function.apply();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
    
    @Override
    public <P> void voidExec(Consumer<P> consumer, P param) {
        try {
            DynamicDataSourceContextHolder.push(datasourceName);
            consumer.accept(param);
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
    
    @Override
    public void voidExec(NoArgConsumer consumer) {
        try {
            DynamicDataSourceContextHolder.push(datasourceName);
            consumer.accept();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
}
